#include "vulkan/vulkan.h"

#ifdef _WIN32
#define NOMINMAX
#include "windows.h"
#include "vulkan/vulkan_win32.h"
#endif // _WIN32

#include "include/core/SkCanvas.h"
#include "include/core/SkColorSpace.h"
#include "include/core/SkImageInfo.h"
#include "include/core/SkRefCnt.h"
#include "include/core/SkSurface.h"
#include "include/core/SkTypes.h"
#include "include/gpu/ganesh/GrTypes.h"
#include "include/gpu/ganesh/GrDirectContext.h"
#include "include/gpu/ganesh/SkSurfaceGanesh.h"
#include "include/gpu/ganesh/vk/GrVkDirectContext.h"
#include "include/gpu/ganesh/vk/GrVkTypes.h"
#include "include/gpu/ganesh/GrBackendSurface.h"  
#include "include/gpu/ganesh/vk/GrVkBackendSurface.h" 
#include "include/gpu/vk/VulkanBackendContext.h"
#include "include/gpu/vk/VulkanExtensions.h"
#include "src/gpu/vk/vulkanmemoryallocator/VulkanMemoryAllocatorPriv.h"
#include "src/gpu/GpuTypesPriv.h"
#include <vector>
#include <cassert>
#include <iostream>

#include "Callback.h"
#include "util/DT.h"

using namespace yzrilyzr_util;

DT refreshDT;

// 全局变量
HINSTANCE hInstance;
HWND hWnd;
VkInstance vkInstance;
VkSurfaceKHR vkSurface;
VkPhysicalDevice vkPhysicalDevice;
VkDevice vkDevice;
VkQueue vkQueue;
uint32_t queueFamilyIndex;
VkSwapchainKHR swapchain=VK_NULL_HANDLE;
std::vector<VkImage> swapchainImages;
VkFormat swapchainImageFormat;
VkExtent2D swapchainExtent;
bool framebufferResized=false;

void cleanupSwapchain();
void recreateSwapchain(sk_sp<GrDirectContext> context, sk_sp<SkSurface> & surface);

LRESULT CALLBACK WndProc(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam){
	switch(msg){
		case WM_DESTROY:
			PostQuitMessage(0);
			return 0;
		case WM_SIZE:
			// 窗口大小改变时标记需要重建交换链
			framebufferResized=true;
			return 0;
		default:
			return DefWindowProc(hWnd, msg, wParam, lParam);
	}
}
void createWindow(){
	WNDCLASS wc={};
	wc.lpfnWndProc=WndProc;
	wc.hInstance=hInstance;
	wc.lpszClassName="SkiaVulkanWindow";
	RegisterClass(&wc);

	hWnd=CreateWindow(
		"SkiaVulkanWindow", "Skia + Vulkan on Windows",
		WS_OVERLAPPEDWINDOW, CW_USEDEFAULT, CW_USEDEFAULT, 800, 600,
		nullptr, nullptr, hInstance, nullptr);
	SetWindowDisplayAffinity(hWnd, WDA_MONITOR);
	ShowWindow(hWnd, SW_SHOW);
}
void cleanupSwapchain(){
	if(swapchain != VK_NULL_HANDLE){
		vkDestroySwapchainKHR(vkDevice, swapchain, nullptr);
		swapchain=VK_NULL_HANDLE;
	}
}
void recreateSwapchain(sk_sp<GrDirectContext> context, sk_sp<SkSurface> & surface){
	// 等待设备空闲
	vkDeviceWaitIdle(vkDevice);

	// 清理旧的交换链
	cleanupSwapchain();

	// 获取新的表面能力
	VkSurfaceCapabilitiesKHR capabilities;
	vkGetPhysicalDeviceSurfaceCapabilitiesKHR(vkPhysicalDevice, vkSurface, &capabilities);

	// 如果窗口最小化了，等待恢复
	if(capabilities.currentExtent.width == 0 || capabilities.currentExtent.height == 0){
		return;
	}

	// 重新创建交换链
	uint32_t formatCount;
	vkGetPhysicalDeviceSurfaceFormatsKHR(vkPhysicalDevice, vkSurface, &formatCount, nullptr);
	std::vector<VkSurfaceFormatKHR> formats(formatCount);
	vkGetPhysicalDeviceSurfaceFormatsKHR(vkPhysicalDevice, vkSurface, &formatCount, formats.data());

	VkSurfaceFormatKHR surfaceFormat=formats[0];
	for(const auto & format : formats){
		if(format.format == VK_FORMAT_B8G8R8A8_UNORM){
			surfaceFormat=format;
			break;
		}
	}

	VkSwapchainCreateInfoKHR swapchainInfo={VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR};
	swapchainInfo.surface=vkSurface;
	swapchainInfo.minImageCount=2;  // Double buffering
	swapchainInfo.imageFormat=surfaceFormat.format;
	swapchainInfo.imageColorSpace=surfaceFormat.colorSpace;
	swapchainInfo.imageExtent=capabilities.currentExtent;
	swapchainInfo.imageArrayLayers=1;
	swapchainInfo.imageUsage=VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;
	swapchainInfo.imageSharingMode=VK_SHARING_MODE_EXCLUSIVE;
	swapchainInfo.preTransform=capabilities.currentTransform;
	swapchainInfo.compositeAlpha=VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
	swapchainInfo.presentMode=VK_PRESENT_MODE_FIFO_KHR;
	swapchainInfo.clipped=VK_TRUE;

	if(vkCreateSwapchainKHR(vkDevice, &swapchainInfo, nullptr, &swapchain) != VK_SUCCESS){
		assert(false && "Failed to recreate swapchain");
	}

	// 获取新的交换链图像
	uint32_t imageCount;
	vkGetSwapchainImagesKHR(vkDevice, swapchain, &imageCount, nullptr);
	swapchainImages.resize(imageCount);
	vkGetSwapchainImagesKHR(vkDevice, swapchain, &imageCount, swapchainImages.data());

	swapchainImageFormat=surfaceFormat.format;
	swapchainExtent=capabilities.currentExtent;

	// 通知窗口大小变化
	onWindowSizeChange(swapchainExtent.width, swapchainExtent.height);

	framebufferResized=false;
}
void initVulkan(){
	// 创建 Vulkan 实例
	VkApplicationInfo appInfo={VK_STRUCTURE_TYPE_APPLICATION_INFO};
	appInfo.pApplicationName="SkiaVulkanApp";
	appInfo.applicationVersion=VK_MAKE_VERSION(1, 0, 0);
	appInfo.pEngineName="No Engine";
	appInfo.engineVersion=VK_MAKE_VERSION(1, 0, 0);
	appInfo.apiVersion=VK_API_VERSION_1_1;

	const char * extensions[]={
		VK_KHR_SURFACE_EXTENSION_NAME,
		VK_KHR_WIN32_SURFACE_EXTENSION_NAME
	};

	VkInstanceCreateInfo instanceInfo={VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO};
	instanceInfo.pApplicationInfo=&appInfo;
	instanceInfo.enabledExtensionCount=2;
	instanceInfo.ppEnabledExtensionNames=extensions;

	if(vkCreateInstance(&instanceInfo, nullptr, &vkInstance) != VK_SUCCESS){
		assert(false && "Failed to create Vulkan instance");
	}

	// 创建 Vulkan 表面
	VkWin32SurfaceCreateInfoKHR surfaceInfo={VK_STRUCTURE_TYPE_WIN32_SURFACE_CREATE_INFO_KHR};
	surfaceInfo.hinstance=hInstance;
	surfaceInfo.hwnd=hWnd;
	if(vkCreateWin32SurfaceKHR(vkInstance, &surfaceInfo, nullptr, &vkSurface) != VK_SUCCESS){
		assert(false && "Failed to create surface");
	}

	// 选择物理设备
	uint32_t physicalDeviceCount=0;
	vkEnumeratePhysicalDevices(vkInstance, &physicalDeviceCount, nullptr);
	std::vector<VkPhysicalDevice> physicalDevices(physicalDeviceCount);
	vkEnumeratePhysicalDevices(vkInstance, &physicalDeviceCount, physicalDevices.data());

	// 选择第一个支持图形和呈现的物理设备
	for(const auto & device : physicalDevices){
		uint32_t queueFamilyCount=0;
		vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, nullptr);
		std::vector<VkQueueFamilyProperties> queueFamilies(queueFamilyCount);
		vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, queueFamilies.data());

		for(uint32_t i=0; i < queueFamilyCount; i++){
			if(queueFamilies[i].queueFlags & VK_QUEUE_GRAPHICS_BIT){
				VkBool32 presentSupport=false;
				vkGetPhysicalDeviceSurfaceSupportKHR(device, i, vkSurface, &presentSupport);
				if(presentSupport){
					vkPhysicalDevice=device;
					queueFamilyIndex=i;
					break;
				}
			}
		}
		if(vkPhysicalDevice != VK_NULL_HANDLE) break;
	}

	if(vkPhysicalDevice == VK_NULL_HANDLE){
		assert(false && "Failed to find suitable physical device");
	}

	// 创建设备和队列
	float queuePriority=1.0f;
	VkDeviceQueueCreateInfo queueInfo={VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO};
	queueInfo.queueFamilyIndex=queueFamilyIndex;
	queueInfo.queueCount=1;
	queueInfo.pQueuePriorities=&queuePriority;

	const char * deviceExtensions[]={VK_KHR_SWAPCHAIN_EXTENSION_NAME};
	VkDeviceCreateInfo deviceInfo={VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO};
	deviceInfo.queueCreateInfoCount=1;
	deviceInfo.pQueueCreateInfos=&queueInfo;
	deviceInfo.enabledExtensionCount=1;
	deviceInfo.ppEnabledExtensionNames=deviceExtensions;

	if(vkCreateDevice(vkPhysicalDevice, &deviceInfo, nullptr, &vkDevice) != VK_SUCCESS){
		assert(false && "Failed to create logical device");
	}
	vkGetDeviceQueue(vkDevice, queueFamilyIndex, 0, &vkQueue);

	// 初始创建交换链
	VkSurfaceCapabilitiesKHR capabilities;
	vkGetPhysicalDeviceSurfaceCapabilitiesKHR(vkPhysicalDevice, vkSurface, &capabilities);

	uint32_t formatCount;
	vkGetPhysicalDeviceSurfaceFormatsKHR(vkPhysicalDevice, vkSurface, &formatCount, nullptr);
	std::vector<VkSurfaceFormatKHR> formats(formatCount);
	vkGetPhysicalDeviceSurfaceFormatsKHR(vkPhysicalDevice, vkSurface, &formatCount, formats.data());

	VkSurfaceFormatKHR surfaceFormat=formats[0];
	for(const auto & format : formats){
		if(format.format == VK_FORMAT_B8G8R8A8_UNORM){
			surfaceFormat=format;
			break;
		}
	}

	VkSwapchainCreateInfoKHR swapchainInfo={VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR};
	swapchainInfo.surface=vkSurface;
	swapchainInfo.minImageCount=2;  // Double buffering
	swapchainInfo.imageFormat=surfaceFormat.format;
	swapchainInfo.imageColorSpace=surfaceFormat.colorSpace;
	swapchainInfo.imageExtent=capabilities.currentExtent;
	swapchainInfo.imageArrayLayers=1;
	swapchainInfo.imageUsage=VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;
	swapchainInfo.imageSharingMode=VK_SHARING_MODE_EXCLUSIVE;
	swapchainInfo.preTransform=capabilities.currentTransform;
	swapchainInfo.compositeAlpha=VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
	swapchainInfo.presentMode=VK_PRESENT_MODE_FIFO_KHR;
	swapchainInfo.clipped=VK_TRUE;

	if(vkCreateSwapchainKHR(vkDevice, &swapchainInfo, nullptr, &swapchain) != VK_SUCCESS){
		assert(false && "Failed to create swapchain");
	}

	// 获取交换链图像
	uint32_t imageCount;
	vkGetSwapchainImagesKHR(vkDevice, swapchain, &imageCount, nullptr);
	swapchainImages.resize(imageCount);
	vkGetSwapchainImagesKHR(vkDevice, swapchain, &imageCount, swapchainImages.data());

	swapchainImageFormat=surfaceFormat.format;
	swapchainExtent=capabilities.currentExtent;
}

sk_sp<GrDirectContext> initSkiaVulkan(){
	skgpu::VulkanBackendContext backendContext;
	backendContext.fInstance=vkInstance;
	backendContext.fPhysicalDevice=vkPhysicalDevice;
	backendContext.fDevice=vkDevice;
	backendContext.fQueue=vkQueue;
	backendContext.fGraphicsQueueIndex=queueFamilyIndex;
	backendContext.fMaxAPIVersion=VK_API_VERSION_1_1;
	backendContext.fVkExtensions=nullptr;
	backendContext.fGetProc=[](const char * proc_name, VkInstance instance, VkDevice device){
		if(device != VK_NULL_HANDLE){
			return vkGetDeviceProcAddr(device, proc_name);
		}
		return vkGetInstanceProcAddr(instance, proc_name);
	};
	backendContext.fProtectedContext=GrProtected::kNo;

	// 初始化内存分配器（关键新增代码）
	backendContext.fMemoryAllocator=skgpu::VulkanMemoryAllocators::Make(
		backendContext, skgpu::ThreadSafe::kNo, std::nullopt);

	return GrDirectContexts::MakeVulkan(backendContext);
}

#include "include/gpu/ganesh/vk/GrVkBackendSurface.h"
#include "include/gpu/ganesh/SkSurfaceGanesh.h"

sk_sp<SkSurface> createSkiaSurface(sk_sp<GrDirectContext> context){
	// 1. 创建VkImageInfo结构体
	GrVkImageInfo imageInfo;
	imageInfo.fImage=swapchainImages[0];
	imageInfo.fAlloc=skgpu::VulkanAlloc();
	imageInfo.fImageTiling=VK_IMAGE_TILING_OPTIMAL;
	imageInfo.fImageLayout=VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
	imageInfo.fFormat=swapchainImageFormat;
	imageInfo.fLevelCount=1;
	imageInfo.fCurrentQueueFamily=queueFamilyIndex;  // 设为实际队列家族索引
	imageInfo.fProtected=GrProtected::kNo;
	imageInfo.fSharingMode=VK_SHARING_MODE_EXCLUSIVE;

	// 2. 创建后端渲染目标
	GrBackendRenderTarget backendRT=GrBackendRenderTargets::MakeVk(
		swapchainExtent.width,
		swapchainExtent.height,
		imageInfo);

	// 3. 创建Skia表面
	SkSurfaceProps props(0, kRGB_H_SkPixelGeometry);

	auto surface=SkSurfaces::WrapBackendRenderTarget(
		context.get(),
		backendRT,
		kTopLeft_GrSurfaceOrigin,
		kBGRA_8888_SkColorType,
		SkColorSpace::MakeSRGB(),  // 使用SRGB颜色空间
		&props);

	if(!surface){
		std::cerr << "Failed to create Skia surface. Backend render target info:\n"
			<< "Width: " << swapchainExtent.width << "\n"
			<< "Height: " << swapchainExtent.height << "\n"
			<< "Format: " << imageInfo.fFormat << "\n";
	}
	onWindowSizeChange(swapchainExtent.width, swapchainExtent.height);
	return surface;
}

void drawFrame(sk_sp<GrDirectContext> context, sk_sp<SkSurface> & surface){
	// 检查是否需要重建交换链
	if(framebufferResized){
		recreateSwapchain(context, surface);
		return; // 跳过这一帧的绘制
	}

	// 获取下一张交换链图像
	uint32_t imageIndex;
	VkSemaphore imageAvailableSemaphore, renderFinishedSemaphore;
	VkFence fence;

	// 创建信号量和栅栏
	VkSemaphoreCreateInfo semaphoreInfo={VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO};
	vkCreateSemaphore(vkDevice, &semaphoreInfo, nullptr, &imageAvailableSemaphore);
	vkCreateSemaphore(vkDevice, &semaphoreInfo, nullptr, &renderFinishedSemaphore);

	VkFenceCreateInfo fenceInfo={VK_STRUCTURE_TYPE_FENCE_CREATE_INFO};
	vkCreateFence(vkDevice, &fenceInfo, nullptr, &fence);

	VkResult result=vkAcquireNextImageKHR(vkDevice, swapchain, UINT64_MAX,
										  imageAvailableSemaphore, VK_NULL_HANDLE, &imageIndex);

					// 检查是否需要重建交换链
	if(result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR){
		framebufferResized=false;
		recreateSwapchain(context, surface);
		return; // 跳过这一帧的绘制
	} else if(result != VK_SUCCESS){
		assert(false && "Failed to acquire swap chain image");
	}

	// 更新图像信息为当前交换链图像
	GrVkImageInfo imageInfo;
	imageInfo.fImage=swapchainImages[imageIndex];
	imageInfo.fAlloc=skgpu::VulkanAlloc();
	imageInfo.fImageTiling=VK_IMAGE_TILING_OPTIMAL;
	imageInfo.fImageLayout=VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
	imageInfo.fFormat=swapchainImageFormat;
	imageInfo.fLevelCount=1;
	imageInfo.fCurrentQueueFamily=VK_QUEUE_FAMILY_IGNORED;
	imageInfo.fProtected=GrProtected::kNo;
	imageInfo.fSharingMode=VK_SHARING_MODE_EXCLUSIVE;

	GrBackendRenderTarget backendRT=GrBackendRenderTargets::MakeVk(
		swapchainExtent.width, swapchainExtent.height, imageInfo);

	// 重新创建表面以使用新的交换链图像
	SkSurfaceProps props(0, kRGB_H_SkPixelGeometry);
	surface=SkSurfaces::WrapBackendRenderTarget(
		context.get(), backendRT, kTopLeft_GrSurfaceOrigin,
		kBGRA_8888_SkColorType, SkColorSpace::MakeSRGB(), &props);

	// 绘制
	SkCanvas * canvas=surface->getCanvas();
	onDraw(*canvas);

	// 提交命令
	context->flushAndSubmit();  // 同步提交

	// 呈现
	VkPresentInfoKHR presentInfo={VK_STRUCTURE_TYPE_PRESENT_INFO_KHR};
	presentInfo.waitSemaphoreCount=1;
	presentInfo.pWaitSemaphores=&renderFinishedSemaphore;
	presentInfo.swapchainCount=1;
	presentInfo.pSwapchains=&swapchain;
	presentInfo.pImageIndices=&imageIndex;

	result=vkQueuePresentKHR(vkQueue, &presentInfo);

	// 检查是否需要重建交换链
	if(result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR || framebufferResized){
		framebufferResized=false;
		recreateSwapchain(context, surface);
	} else if(result != VK_SUCCESS){
		assert(false && "Failed to present swap chain image");
	}

	// 等待GPU完成工作
	vkQueueWaitIdle(vkQueue);

	// 清理
	vkDestroySemaphore(vkDevice, imageAvailableSemaphore, nullptr);
	vkDestroySemaphore(vkDevice, renderFinishedSemaphore, nullptr);
	vkDestroyFence(vkDevice, fence, nullptr);
}

int WINAPI WinMain(HINSTANCE hInst, HINSTANCE, LPSTR, int){
	hInstance=hInst;
	createWindow();
	initVulkan();

	auto context=initSkiaVulkan();
	if(!context){
		assert(false && "Failed to create Skia Vulkan context");
		return -1;
	}
	onCreateApp();

	// 创建与交换链关联的Skia表面
	auto surface=createSkiaSurface(context);
	if(!surface){
		assert(false && "Failed to create Skia surface");
		return -1;
	}

	MSG msg={};
	const UINT64 frameTime=16666; // 约60 FPS (1秒/60 ≈ 16.666毫秒，单位为100纳秒)
	LARGE_INTEGER frequency, lastTime;
	QueryPerformanceFrequency(&frequency);
	QueryPerformanceCounter(&lastTime);
	while(true){
		// 处理窗口消息，但不阻塞等待
		if(PeekMessage(&msg, nullptr, 0, 0, PM_REMOVE)){
			if(msg.message == WM_QUIT){
				break; // 收到退出消息时终止循环
			}
			TranslateMessage(&msg);
			DispatchMessage(&msg);
		}

		// 帧率控制：确保每帧间隔约16ms
		LARGE_INTEGER currentTime;
		QueryPerformanceCounter(&currentTime);
		UINT64 elapsedTime=(currentTime.QuadPart - lastTime.QuadPart) * 1000000000 / frequency.QuadPart;

		if(elapsedTime >= frameTime){
			drawFrame(context, surface); // 绘制新帧
			lastTime=currentTime;      // 更新时间戳
		} else{
			// 短暂休眠，减少CPU占用
			Sleep(1);
		}
	}


	// 清理资源
	vkDestroySwapchainKHR(vkDevice, swapchain, nullptr);
	vkDestroyDevice(vkDevice, nullptr);
	vkDestroySurfaceKHR(vkInstance, vkSurface, nullptr);
	vkDestroyInstance(vkInstance, nullptr);

	return 0;
}
int main(){
	return WinMain(GetModuleHandle(NULL), NULL, NULL, SW_SHOW);
}